import os
from datetime import datetime, timedelta

import numpy as np
import pandas as pd

# 设置随机种子以确保结果可重现
np.random.seed(42)

# 创建一个包含各种问题列的样本数据集
def create_sample_dataset(rows=1000):
    """创建包含各种无效列类型的样本数据集"""
    
    # 正常列：数值
    normal_numeric = np.random.normal(50, 10, rows)
    
    # 正常列：分类
    categories = ['类别A', '类别B', '类别C', '类别D', '类别E']
    normal_categorical = np.random.choice(categories, size=rows)
    
    # 正常列：日期（过去30天内）
    today = datetime.now()
    normal_dates = [today - timedelta(days=np.random.randint(0, 30)) for _ in range(rows)]
    
    # 问题列1：高缺失值
    high_missing = normal_numeric.copy()
    missing_mask = np.random.random(rows) < 0.6  # 60% 的值会被设为NaN
    high_missing[missing_mask] = np.nan
    
    # 问题列2：常数值
    constant_value = np.full(rows, 42)
    
    # 问题列3：高基数（唯一值）
    high_cardinality = [f"唯一ID_{i}" for i in range(rows)]
    
    # 问题列4：低方差
    low_variance = np.random.normal(100, 0.01, rows)  # 均值100，标准差0.01
    
    # 问题列5：高相关性
    high_correlation = normal_numeric * 2 + 5 + np.random.normal(0, 0.1, rows)  # 与normal_numeric高度相关
    
    # 问题列6：空字符串问题
    string_issues = []
    for _ in range(rows):
        r = np.random.random()
        if r < 0.4:  # 40% 空字符串
            string_issues.append("")
        elif r < 0.6:  # 20% 空格字符串
            string_issues.append("   ")
        else:  # 40% 正常字符串
            string_issues.append(f"正常文本-{np.random.randint(1, 100)}")
    
    # 问题列7：ID列（可能导致数据泄漏）
    id_column = [f"ID{i:06d}" for i in range(rows)]
    
    # 问题列8：未来日期
    future_dates = [today + timedelta(days=np.random.randint(1, 365)) for _ in range(rows)]
    
    # 创建DataFrame
    df = pd.DataFrame({
        # 正常列
        'normal_numeric': normal_numeric,
        'normal_categorical': normal_categorical,
        'normal_date': normal_dates,
        
        # 问题列
        'high_missing_values': high_missing,
        'constant_column': constant_value,
        'unique_id': high_cardinality,
        'low_variance_column': low_variance,
        'correlated_column': high_correlation,
        'string_with_issues': string_issues,
        'customer_id': id_column,
        'future_date': future_dates
    })
    
    return df

# 创建并保存样本数据集
if __name__ == "__main__":
    # 创建数据
    df = create_sample_dataset(1000)
    
    # 确保输出目录存在
    output_dir = "sample_data"
    os.makedirs(output_dir, exist_ok=True)
    
    # 保存为CSV和XLSX
    csv_path = os.path.join(output_dir, "sample_data_with_issues.csv")
    xlsx_path = os.path.join(output_dir, "sample_data_with_issues.xlsx")
    
    df.to_csv(csv_path, index=False, encoding='utf-8')
    df.to_excel(xlsx_path, index=False)
    
    print("样本数据已保存到：")
    print(f"CSV: {csv_path}")
    print(f"XLSX: {xlsx_path}")
    
    # 显示数据概览
    print("\n数据概览：")
    print(df.head())
    
    print("\n数据信息：")
    print(df.info())
    
    print("\n预期的无效列：")
    print("- high_missing_values: 高比例缺失值")
    print("- constant_column: 常数值列")
    print("- unique_id: 高基数列（几乎每行都是唯一值）")
    print("- low_variance_column: 低方差列")
    print("- correlated_column: 高相关性列（与normal_numeric高度相关）")
    print("- string_with_issues: 字符串问题列（高比例空字符串和空格）")
    print("- customer_id: 潜在的数据泄漏列（ID列）")
    print("- future_date: 日期问题列（未来日期）")
